#pragma once

#include <memory>
#include <atomic>
#include <compare>
#include <utility>
#include <concepts>
#ifndef NDEBUG
#include <cstring>
#endif
#include <zenox/pragma.hpp>

namespace zenox {

namespace _memory_details {

struct alignas(__STDCPP_DEFAULT_NEW_ALIGNMENT__) _shared_block {
    size_t m_share_cnt;
    size_t m_weak_cnt;
};

struct _tag_private_construct {
    explicit _tag_private_construct() = default;
};

struct _tag_static_pointer_cast {
    explicit _tag_static_pointer_cast() = default;
};

struct _tag_dynamic_pointer_cast {
    explicit _tag_dynamic_pointer_cast() = default;
};

}

template <class T>
struct weak_ptr;

ZENO_GNUC_ATTRIBUTE(malloc, nonnull) inline void *_allocate(size_t n) {
    return ::operator new(n);
}

ZENO_GNUC_ATTRIBUTE(nonnull(1)) inline void _deallocate(void *p) noexcept {
    ::operator delete(p);
}

template <class T>
struct shared_ptr {
    static_assert(std::is_same_v<std::decay_t<T>, T>, "we currently only support cv-unqualified type for now");

private:
    T *m_ptr;
    _memory_details::_shared_block *_block() const noexcept {
#ifndef NDEBUG
        return m_ptr ? reinterpret_cast<_memory_details::_shared_block *>(m_ptr) - 1 : nullptr;
#else
        return reinterpret_cast<_memory_details::_shared_block *>(m_ptr) - 1;
#endif
    }

    template <class U>
    friend struct shared_ptr;

    template <class U>
    friend struct weak_ptr;

public:
    explicit shared_ptr(_memory_details::_tag_private_construct, T *ptr) noexcept
    : m_ptr(ptr)
    {}

    shared_ptr() noexcept
    : m_ptr(nullptr)
    {}

    shared_ptr(std::nullptr_t) noexcept
    : m_ptr(nullptr)
    {}

    template <class ...Ts>
    explicit shared_ptr(std::in_place_t, Ts &&...ts) {
        static_assert(alignof(T) <= __STDCPP_DEFAULT_NEW_ALIGNMENT__,
            "we can't handle type with alignment larger than __STDCPP_DEFAULT_NEW_ALIGNMENT__");
        m_ptr = reinterpret_cast<T *>(reinterpret_cast<_memory_details::_shared_block *>(
            _allocate(sizeof(T) + sizeof(_memory_details::_shared_block))) + 1);
        std::construct_at(m_ptr, std::forward<Ts>(ts)...);
        _block()->m_share_cnt = 1;
        _block()->m_weak_cnt = 1;
    }

    
    explicit shared_ptr(_memory_details::_tag_private_construct, std::in_place_t)
    requires std::default_initializable<T>
    {
        static_assert(alignof(T) <= __STDCPP_DEFAULT_NEW_ALIGNMENT__,
            "we can't handle type with alignment larger than __STDCPP_DEFAULT_NEW_ALIGNMENT__");
        m_ptr = reinterpret_cast<T *>(reinterpret_cast<_memory_details::_shared_block *>(
            _allocate(sizeof(T) + sizeof(_memory_details::_shared_block))) + 1);
        _block()->m_share_cnt = 1;
        _block()->m_weak_cnt = 1;
    }

    template <std::derived_from<T> U> requires (!std::same_as<U, T>)
    shared_ptr(shared_ptr<U> that) noexcept
    : m_ptr(std::exchange(that.m_ptr, nullptr))
    {}

    template <class U>
    explicit shared_ptr(_memory_details::_tag_static_pointer_cast, shared_ptr<U> that) noexcept
    : m_ptr(static_cast<T *>(std::exchange(that.m_ptr, nullptr)))
    {}

    template <class U>
    explicit shared_ptr(_memory_details::_tag_dynamic_pointer_cast, shared_ptr<U> that) noexcept
    : m_ptr(dynamic_cast<T *>(std::exchange(that.m_ptr, nullptr)))
    {}

    shared_ptr(shared_ptr &&that) noexcept
    : m_ptr(std::exchange(that.m_ptr, nullptr))
    {}

    shared_ptr(shared_ptr const &that) noexcept
    : m_ptr(that.m_ptr)
    {
        if (m_ptr) {
            ++_block()->m_share_cnt;
            ++_block()->m_weak_cnt;
        }
    }

    shared_ptr &operator=(shared_ptr &&that) noexcept {
        if (this == &that) [[unlikely]] return *this;
        if (m_ptr) {
            if (!--_block()->m_share_cnt) {
                m_ptr->~T();
#ifndef NDEBUG
                std::memset((void *)m_ptr, 0xcc, sizeof(T));
#endif
            }
            if (!--_block()->m_weak_cnt) {
                _deallocate(_block());
            }
        }
        m_ptr = std::exchange(that.m_ptr, nullptr);
        return *this;
    }

    shared_ptr &operator=(shared_ptr const &that) noexcept {
        if (this == &that) [[unlikely]] return *this;
        if (m_ptr) {
            if (!--_block()->m_share_cnt) {
                m_ptr->~T();
#ifndef NDEBUG
                std::memset((void *)m_ptr, 0xcc, sizeof(T));
#endif
            }
            if (!--_block()->m_weak_cnt) {
                _deallocate(_block());
            }
        }
        m_ptr = that.m_ptr;
        if (m_ptr) {
            ++_block()->m_share_cnt;
            ++_block()->m_weak_cnt;
        }
        return *this;
    }

    ~shared_ptr() noexcept {
        if (m_ptr) {
            if (!--_block()->m_share_cnt) {
                m_ptr->~T();
#ifndef NDEBUG
                std::memset((void *)m_ptr, 0xcc, sizeof(T));
#endif
            }
            if (!--_block()->m_weak_cnt) {
                _deallocate(_block());
            }
        }
    }

    size_t use_count() const noexcept {
        return m_ptr ? _block()->m_share_cnt : 0;
    }

    void reset() noexcept {
        if (m_ptr) {
            if (!--_block()->m_share_cnt) {
                m_ptr->~T();
#ifndef NDEBUG
                std::memset((void *)m_ptr, 0xcc, sizeof(T));
#endif
            }
            if (!--_block()->m_weak_cnt) {
                _deallocate(_block());
            }
        }
        m_ptr = nullptr;
        _block() = nullptr;
    }

    T *get() const noexcept {
        return m_ptr;
    }

    weak_ptr<T> weak() const noexcept {
        return weak_ptr<T>(*this);
    }

    T *operator->() const noexcept {
        return m_ptr;
    }

    explicit operator bool() const noexcept {
        return m_ptr;
    }

    bool operator==(shared_ptr<T> const &that) const noexcept {
        return m_ptr == that.m_ptr;
    }

    bool operator==(weak_ptr<T> const &that) const noexcept {
        return m_ptr == that.m_ptr;
    }

    bool operator==(T *that) const noexcept {
        return m_ptr == that;
    }

    std::strong_ordering operator<=>(shared_ptr<T> const &that) const noexcept {
        return m_ptr <=> that.m_ptr;
    }

    std::strong_ordering operator<=>(weak_ptr<T> const &that) const noexcept {
        return m_ptr <=> that.m_ptr;
    }

    std::strong_ordering operator<=>(T *that) const noexcept {
        return m_ptr <=> that;
    }
};

template <class T>
struct weak_ptr {
    static_assert(std::is_same_v<std::decay_t<T>, T>, "we currently only support cv-unqualified type for now");

private:
    T *m_ptr;
    _memory_details::_shared_block *_block() const noexcept {
#ifndef NDEBUG
        return m_ptr ? reinterpret_cast<_memory_details::_shared_block *>(m_ptr) - 1 : nullptr;
#else
        return reinterpret_cast<_memory_details::_shared_block *>(m_ptr) - 1;
#endif
    }

    template <class U>
    friend struct shared_ptr;

    template <class U>
    friend struct weak_ptr;

public:
    explicit weak_ptr(_memory_details::_tag_private_construct, T *ptr) noexcept
    : m_ptr(ptr)
    {}

    weak_ptr() noexcept
    : m_ptr(nullptr)
    {}

    weak_ptr(std::nullptr_t) noexcept
    : m_ptr(nullptr)
    {}

    weak_ptr(shared_ptr<T> const &sp) noexcept
    : m_ptr(sp.m_ptr)
    {
        if (m_ptr)
            ++_block()->m_weak_cnt;
    }

    weak_ptr(weak_ptr &&that) noexcept
    : m_ptr(std::exchange(that.m_ptr, nullptr))
    {}

    weak_ptr(weak_ptr const &that) noexcept
    : m_ptr(that.m_ptr)
    {
        if (m_ptr)
            ++_block()->m_weak_cnt;
    }

    template <std::derived_from<T> U> requires (!std::same_as<U, T>)
    weak_ptr(weak_ptr<U> that) noexcept
    : m_ptr(std::exchange(that.m_ptr, nullptr))
    {}

    weak_ptr &operator=(weak_ptr &&that) noexcept {
        if (this == &that) [[unlikely]] return *this;
        if (m_ptr && !--_block()->m_weak_cnt) {
            _deallocate(_block());
        }
        m_ptr = std::exchange(that.m_ptr, nullptr);
        return *this;
    }

    weak_ptr &operator=(weak_ptr const &that) noexcept {
        if (this == &that) [[unlikely]] return *this;
        if (m_ptr && !--_block()->m_weak_cnt) {
            _deallocate(_block());
        }
        m_ptr = that.m_ptr;
        if (m_ptr) ++_block()->m_weak_cnt;
        return *this;
    }

    ~weak_ptr() noexcept {
        if (m_ptr && !--_block()->m_weak_cnt) {
            _deallocate(_block());
        }
    }

    void reset() noexcept {
        m_ptr = nullptr;
        if (m_ptr && !--_block()->m_weak_cnt) {
            _deallocate(_block());
        }
        _block() = nullptr;
    }

    size_t use_count() const noexcept {
        return m_ptr ? _block()->m_share_cnt : 0;
    }

    bool expired() const noexcept {
        return !(m_ptr && _block()->m_share_cnt);
    }

    shared_ptr<T> lock() const noexcept {
        if (m_ptr && _block()->m_share_cnt) {
            ++_block()->m_share_cnt;
            ++_block()->m_weak_cnt;
            return shared_ptr<T>(_memory_details::_tag_private_construct(), m_ptr);
        } else {
            return nullptr;
        }
    }

    T *get() const noexcept {
        if (m_ptr && _block()->m_share_cnt) {
            return m_ptr;
        } else {
            return nullptr;
        }
    }

    T *get_unsafe() const noexcept {
        return m_ptr;
    }

    T *operator->() const noexcept {
        return get();
    }

    explicit operator bool() const noexcept {
        return m_ptr && _block()->m_share_cnt;
    }

    bool operator==(shared_ptr<T> const &that) const noexcept {
        return m_ptr == that.m_ptr;
    }

    bool operator==(weak_ptr<T> const &that) const noexcept {
        return m_ptr == that.m_ptr;
    }

    bool operator==(T *that) const noexcept {
        return m_ptr == that;
    }

    std::strong_ordering operator<=>(shared_ptr<T> const &that) const noexcept {
        return m_ptr <=> that.m_ptr;
    }

    std::strong_ordering operator<=>(weak_ptr<T> const &that) const noexcept {
        return m_ptr <=> that.m_ptr;
    }

    std::strong_ordering operator<=>(T *that) const noexcept {
        return m_ptr <=> that;
    }
};

template <class T>
inline weak_ptr<T> weak_from(T *p) { // usage: weak_from(this)
    auto blk = reinterpret_cast<_memory_details::_shared_block *>(p) - 1;
    ++blk->m_weak_cnt;
    return weak_ptr<T>(_memory_details::_tag_private_construct(), p);
}

template <class T>
inline shared_ptr<T> shared_from(T *p) { // usage: shared_from(this)
    auto blk = reinterpret_cast<_memory_details::_shared_block *>(p) - 1;
    ++blk->m_share_cnt;
    ++blk->m_weak_cnt;
    return shared_ptr<T>(_memory_details::_tag_private_construct(), p);
}

template <class T, class ...Ts> requires std::constructible_from<T, Ts...>
inline shared_ptr<T> make_shared(Ts &&...ts) {
    return shared_ptr<T>(std::in_place, std::forward<Ts>(ts)...);
}

template <std::default_initializable T>
inline shared_ptr<T> make_shared_for_overwrite() {
    return shared_ptr<T>(_memory_details::_tag_private_construct(), std::in_place);
}

template <class To, class T>
inline shared_ptr<To> static_pointer_cast(shared_ptr<T> p) {
    return shared_ptr<To>(_memory_details::_tag_static_pointer_cast(), std::move(p));
}

template <class To, class T>
inline shared_ptr<To> dynamic_pointer_cast(shared_ptr<T> p) {
    return shared_ptr<To>(_memory_details::_tag_dynamic_pointer_cast(), std::move(p));
}

}
